#include "CoderToJpg.h"

CoderToJpg::CoderToJpg()
:m_nWidth(2592)
,m_nHeight(1520)
{

}

CoderToJpg::~CoderToJpg()
{

}

void CoderToJpg::setWidth(const int p_nWidth)
{
	m_nWidth = p_nWidth;
}

void CoderToJpg::setHeight(const int p_nHeight)
{
	m_nHeight = p_nHeight;
}

int CoderToJpg::init()
{
	av_register_all();
	m_pobjJpgAVFormatContext = avformat_alloc_context();
	if (m_pobjJpgAVFormatContext == NULL)
	{
		return -1;
	}

	m_pobjAVOutputFormat = av_guess_format("mjpeg", NULL, NULL);
	if (m_pobjAVOutputFormat == NULL)
	{
		return -1;
	}

	m_pobjJpgAVFormatContext->oformat = m_pobjAVOutputFormat;

	//
	m_pobjAVStream = avformat_new_stream(m_pobjJpgAVFormatContext, NULL);
	if (m_pobjAVStream == NULL)
	{
		return -1;
	}

	m_pobjJpgAVCodecCtx = m_pobjAVStream->codec;
	m_pobjJpgAVCodecCtx->codec_id = m_pobjAVOutputFormat->video_codec;
	m_pobjJpgAVCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
	m_pobjJpgAVCodecCtx->pix_fmt = AV_PIX_FMT_YUVJ420P;
	m_pobjJpgAVCodecCtx->width = m_nWidth;
	m_pobjJpgAVCodecCtx->height = m_nHeight;
	m_pobjJpgAVCodecCtx->time_base.num = 1;
	m_pobjJpgAVCodecCtx->time_base.den = 25;

	m_pobjJpgAVCodec = avcodec_find_encoder(m_pobjJpgAVCodecCtx->codec_id);
	if (!m_pobjJpgAVCodec)
	{
		return -1;
	}

	int nRet = avcodec_open2(m_pobjJpgAVCodecCtx, m_pobjJpgAVCodec, NULL);
	if (nRet < 0)
	{
		printf("Could not open codec.");
		return -1;
	}

	int y_size = m_nWidth * m_nHeight;
	m_pobjJpgAVPacket = (AVPacket *)malloc(sizeof(AVPacket));
	av_new_packet(m_pobjJpgAVPacket, y_size * 3);

	return 0;
}

int CoderToJpg::coderFrameToPic(AVFrame* p_pobjAVFrame, const string& p_strJpgFileName)
{
	if (p_pobjAVFrame == NULL)
	{
		return -1;
	}

	int nRet = avio_open(&m_pobjJpgAVFormatContext->pb, p_strJpgFileName.c_str(), AVIO_FLAG_READ_WRITE);
	if (nRet< 0)
	{
		return -1;
	}

	//Write Header
	avformat_write_header(m_pobjJpgAVFormatContext, NULL);

	//Encode
	int got_picture = -1;
	nRet = avcodec_encode_video2(m_pobjJpgAVCodecCtx, m_pobjJpgAVPacket, p_pobjAVFrame, &got_picture);
	if(nRet < 0)
	{
		printf("Encode Error.\n");
		return -1;
	}

	if (got_picture == 1)
	{
		m_pobjJpgAVPacket->stream_index = m_pobjAVStream->index;
		nRet = av_write_frame(m_pobjJpgAVFormatContext, m_pobjJpgAVPacket);
	}

	av_free_packet(m_pobjJpgAVPacket);

	//Write Trailer  
	av_write_trailer(m_pobjJpgAVFormatContext);
	avio_close(m_pobjJpgAVFormatContext->pb);
	return 0;
}

void CoderToJpg::close()
{
	avcodec_close(m_pobjAVStream->codec);
}
